//QT api
#include <QString>
#include <QtCore>
#include <QDir>
#include <QDebug>
#include <QApplication>
//OpenCV api
#include "opencv2/features2d/features2d.hpp"
#include "opencv2/nonfree/nonfree.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/ml/ml.hpp"
//BOOST
#include <boost/typeof/typeof.hpp>
//C++ api
#include <ctime>

using namespace std;
using namespace cv;
// some utility functions
bool MakeDir( const string& filepath );
void help( const char* progName );
void GetDirList( string directory, vector<string>& dirlist );
void GetFileList( string directory, vector<string>& filelist );

const string kVocabularyFile("vocabulary.xml.gz");
const string kBowImageDescriptorsDir("/bagOfWords");
const string kSvmsDirs("/svms");

class Params {
public:
    Params(): wordCount( 1000 ), detectorType( "SURF" ),
        descriptorType( "SURF" ), matcherType( "FlannBased" ){ }
    int		wordCount;
    string	detectorType;
    string	descriptorType;
    string	matcherType;
};

/*
 * loop through every directory
 * compute each image's keypoints and descriptors
 * train a vocabulary
 */
Mat BuildVocabulary( const string& databaseDir,
                     const vector<string>& categories,
                     const Ptr<FeatureDetector>& detector,
                     const Ptr<DescriptorExtractor>& extractor,
                     int wordCount) {
    Mat allDescriptors;
    for ( int index = 0; index != (int)categories.size(); ++index ) {
        cout << "---processing category " << categories[index] <<"---"<< endl;
        string currentCategory = databaseDir + "/" + categories[index];
        vector<string> filelist;
        GetFileList( currentCategory, filelist);
        for (BOOST_AUTO(fileindex,filelist.begin()); fileindex != filelist.end(); ++fileindex ) {
            string filepath = currentCategory + "/" + *fileindex;
            cout<<"***processing file "<<filepath<<"***"<<endl;
            Mat image = imread( filepath );
            if (image.empty()) {
                cout<<*fileindex<<" is empty."<<endl;
                continue; // maybe not an image file
            }
            vector<KeyPoint> keyPoints;
            Mat descriptors;
            detector -> detect( image, keyPoints);
            extractor -> compute( image, keyPoints, descriptors );
            if ( allDescriptors.empty() )
                allDescriptors.create( 0, descriptors.cols, descriptors.type() );
            allDescriptors.push_back( descriptors );
        }
        cout << "###done processing category " << categories[index]<<" ###"<< endl;
    }
    assert( !allDescriptors.empty() );
    cout<<"allDescriptors.size="<<allDescriptors.size()<<endl;
    cout << "build vocabulary..." << endl;
    BOWKMeansTrainer bowTrainer( wordCount );
    Mat vocabulary = bowTrainer.cluster( allDescriptors );
    cout << "done build vocabulary..." << endl;
    return vocabulary;
}

// bag of words of an image as its descriptor, not keypoint descriptors
void ComputeBowImageDescriptors( const string& databaseDir,
                                 const vector<string>& categories,
                                 const Ptr<FeatureDetector>& detector,
                                 Ptr<BOWImgDescriptorExtractor>& bowExtractor,
                                 const string& imageDescriptorsDir,
                                 map<string, Mat>* samples) {
    for (int i = 0; i != (int)categories.size(); ++i ) {
        string currentCategory = databaseDir + "/" + categories[i];
        vector<string> filelist;
        GetFileList( currentCategory, filelist);
        for (BOOST_AUTO(fileitr,filelist.begin()); fileitr != filelist.end(); ++fileitr ) {
            string descriptorFileName = imageDescriptorsDir + "/" + (*fileitr) + ".xml.gz";
            FileStorage fs( descriptorFileName, FileStorage::READ );
            Mat imageDescriptor;
            if ( fs.isOpened() ) { // already cached
                fs["imageDescriptor"] >> imageDescriptor;
            } else {
                string filepath = currentCategory + "/" + *fileitr;
                Mat image = imread( filepath );
                if ( image.empty() ) {
                    continue; // maybe not an image file
                }
                vector<KeyPoint> keyPoints;
                detector -> detect( image, keyPoints );
                bowExtractor -> compute( image, keyPoints, imageDescriptor );
                fs.open( descriptorFileName, FileStorage::WRITE );
                if ( fs.isOpened() )
                    fs << "imageDescriptor" << imageDescriptor;
                else
                    cout<<"cannot open file "<<descriptorFileName<<endl;
            }
            if ( samples -> count( categories[i] ) == 0 ) {
                ( *samples )[categories[i]].create( 0, imageDescriptor.cols, imageDescriptor.type() );
            }
            ( *samples )[categories[i]].push_back( imageDescriptor );
        }
    }
}

void train( const map<string, Mat>& samples, const string& category, const CvSVMParams& svmParams, CvSVM* svm ) {
    Mat allSamples( 0, samples.at( category ).cols, samples.at( category ).type() );
    Mat responses( 0, 1, CV_32SC1 );
    //assert( responses.type() == CV_32SC1 );
    allSamples.push_back( samples.at( category ) );
    Mat posResponses( samples.at( category ).rows, 1, CV_32SC1, Scalar::all(1) );
    responses.push_back( posResponses );

    for (BOOST_AUTO(itr,samples.begin()); itr != samples.end(); ++itr ) {
        if ( itr -> first == category ) {
            continue;
        }
        allSamples.push_back( itr -> second );
        Mat response( itr -> second.rows, 1, CV_32SC1, Scalar::all( -1 ) );
        responses.push_back( response );
    }
    svm -> train( allSamples, responses, Mat(), Mat(), svmParams );
}

void trainSVM(const map<string, Mat>& samples, const string& svmDir)
{
    SVMParams svmParams;
    int sign = 0; //sign of the positive class
    float confidence = -FLT_MAX;
    for (BOOST_AUTO(itr,samples.begin()); itr != samples.end(); ++itr ) {
        CvSVM svm;
        string svmFileName = svmDir + "/" + itr -> first + ".xml.gz";
        FileStorage fs( svmFileName, FileStorage::READ );
        if ( fs.isOpened() ) { // exist a previously trained svm
            fs.release();
            svm.load( svmFileName.c_str() );
        } else {
            train( samples, itr->first, svmParams, &svm );
            if ( !svmDir.empty() ) {
                svm.save( svmFileName.c_str() );
            }
        }
    }
}

// using 1-vs-all method, train an svm for each category.
// choose the category with the biggest confidence
string ClassifyBySvm( const Mat& queryDescriptor, const map<string, Mat>& samples, const string& svmDir ) {
    string category;
    int sign = 0; //sign of the positive class
    float confidence = -FLT_MAX;

    for (BOOST_AUTO(itr,samples.begin()); itr != samples.end(); ++itr ) {
        CvSVM svm;
        string svmFileName = svmDir + "\\" + itr -> first + ".xml.gz";
        FileStorage fs( svmFileName, FileStorage::READ );
        if (!fs.isOpened() )
             return "NO_SVM_LOADED";
        fs.release();   // exist a previously trained svm
        svm.load( svmFileName.c_str() );
        // determine the sign of the positive class
        if ( sign == 0 ) {
            float scoreValue = svm.predict( queryDescriptor, true );
            float classValue = svm.predict( queryDescriptor, false );
            sign = ( scoreValue < 0.0f ) == ( classValue < 0.0f )? 1 : -1;
        }
        float curConfidence = sign * svm.predict( queryDescriptor, true );
        if ( curConfidence > confidence ) {
            confidence = curConfidence;
            category = itr -> first;
        }
    }
    return category;
}
//从其他抽象图像类中随机选择(100/num)%幅图像进行当前抽象类视觉词典的计算,
void randomSelectImage(string databaseDir, string currentAbsCategoryName, string destinationDir)
{
    //首先根据抽象及其具体类别数确定从每个具体类中选多少张图像
    vector<string> absCategories;
    GetDirList(databaseDir, absCategories);    //获取抽象类名称
    float percent =  (int)databaseDir.size() / 100.0; //每个抽象类的百分比

    int seed = 0;
    for(int i = 0; i < (int)absCategories.size(); ++i) {
        string absPath = databaseDir + "/" + absCategories[i];
        vector<string> inheritedCategories;
        GetDirList(absPath, inheritedCategories);
        MakeDir(destinationDir + "/" + absCategories[i]);

        for(int j = 0; j < (int)inheritedCategories.size(); ++j) {
            string inhPath = absPath + "/" + inheritedCategories[j];
            vector<string> inheritedCategoryFileName;
            GetFileList(inhPath,inheritedCategoryFileName);
            //从每个具体类别中选择percent%的图像，无视图像数
            int numberToSelect = (int)inheritedCategoryFileName.size() * percent;

            //            MakeDir(destinationDir + "/" + absCategories[i] + "/" + inheritedCategories[j]);
            for(int k = 0; k < numberToSelect; ++k) {  //XXX 为所需选择的图像数
                srand(time(0) + (++seed));
                int random_num = rand() % (int)inheritedCategoryFileName.size();
                QFile file(QString::fromStdString(inhPath + "/" + inheritedCategoryFileName[k]));
                if(file.exists()) {
                    if(!file.copy(QString::fromStdString(destinationDir + "/" +
                                                         absCategories[i] + "/" +
                                                         inheritedCategories[j] + "_" + //重名问题很麻烦-_-！
                                                         inheritedCategoryFileName[random_num])))
                        cout<<"###E### Copy file "<<inheritedCategoryFileName[k]
                              <<" from "<<inhPath
                             <<"to "<<destinationDir<<" failed."<<endl;
                }
            }
        }
    }
}

map<string, Mat> bow(Params &params,string vocabularyFile, Mat &vocabulary, string categoriesPath, string bowIhrImageDescriptorDir,
                     vector<string> categories)
{
    Ptr<FeatureDetector> detector = FeatureDetector::create( params.descriptorType );
    Ptr<DescriptorExtractor> extractor = DescriptorExtractor::create( params.descriptorType );
    Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create( params.matcherType );

    FileStorage fs( vocabularyFile, FileStorage::READ );
    if ( fs.isOpened() ) {
        fs["vocabulary"] >> vocabulary;
    } else {
        vocabulary = BuildVocabulary(categoriesPath,
                                     categories,
                                     detector, extractor, params.wordCount);
        FileStorage fs( vocabularyFile, FileStorage::WRITE );
        if ( fs.isOpened() )
            fs << "vocabulary" << vocabulary;
        else
            cout<<"###E### Cannot save vocabulary file!"<<endl;
    }
    Ptr<BOWImgDescriptorExtractor> bowExtractor = new BOWImgDescriptorExtractor( extractor, matcher );
    bowExtractor -> setVocabulary( vocabulary );
    map<string, Mat> samples;//key: category name, value: histogram
    cout<<"Computing BoW image descriptors"<<endl;
    ComputeBowImageDescriptors(categoriesPath,
                               categories,
                               detector, bowExtractor, bowIhrImageDescriptorDir,  &samples );
    cout<<"###I### Done computing BoW image descriptos"<<endl;
    return samples;
}

string classifyAbsCategory(Params &params,)
{
    Ptr<FeatureDetector> detector = FeatureDetector::create( params.descriptorType );
    Ptr<DescriptorExtractor> extractor = DescriptorExtractor::create( params.descriptorType );
    Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create( params.matcherType );

    Mat image = imread(currentTestDir + "/" + testFileName[j]);
    vector<KeyPoint> keyPoints;
    detector -> detect( image, keyPoints );
    Ptr<BOWImgDescriptorExtractor> bowExtractor = new BOWImgDescriptorExtractor( extractor, matcher );
    bowExtractor -> setVocabulary( vocabulary );
    Mat queryDescriptor;
    bowExtractor -> compute( image, keyPoints, queryDescriptor );
    string category = ClassifyBySvm( queryDescriptor, samples, svmsDir );   //method == svm
    if(category == testCategories[i])
        ++correctTests;
}

//先分抽象类，再分具体类
void test(Params &params, string testDir)
{
    vector<string> testAbsCategories;
    GetDirList(testDir,testAbsCategories);
    for(int i = 0; i < (int)testAbsCategories.size(); ++i) {
        int accurateAbsNum = 0, accurateInhNum = 0;
        cout<<"###I### Testing category "<<testAbsCategories[i]<<endl;
        string currentInhTestDir = testDir + "/" +  testAbsCategories[i];
        vector<string> testFileName;
        GetDirList(currentInhTestDir,testFileName);
        for(int j = 0; j < (int)currentInhTestDir.size(); ++j) {
            int correctTests = 0;
            for(int k = 0; k < (int)testFileName.size(); ++k) {
                cout<<"###I### File progress: "<<k<<"/"<<(int)testFileName.size()<<
                   " Inherited categroy progress: "<<j<<"/"<<(int)currentInhTestDir.size()<<
                      " Abstracted category progress: "<<i<<"/"<<(int)testAbsCategories.size()<<endl;
                string absCategory = classifyAbsCategory(params,);
                string inhCategory = classifyInhCategory(params,absCategory);
                if(inhCategory == currentInhTestDir[j])
                    ++accurateInhNum;
                if(absCategory == testAbsCategories[i])
                    ++accurateAbsNum;
            }
        }
        cout<<"###I### Testing Category "<<testAbsCategories[i]<<" finished."<<endl;
        cout<<"Accurate Rate: "<<endl;
        cout<<"Abstracted category: "<<(float)accurateAbsNum / (float)testFileName.size()<<endl;
        cout<<"Inherited category: "<<(float)accurateInhNum / (float)testFileName.size()<<endl;
    }
}

void classifyInhCategory(Params &params, string testDir)
{

}

int main( int argc, char* argv[] ) {
    QCoreApplication a(argc, argv);
    if ( argc != 5 && argc != 8 ) {
        help( argv[0] );
        return -1;
    }
    cv::initModule_nonfree();
    // read params
    Params params;
    string method = argv[1];    //svm
    string testDir = argv[2];    //测试图像
    string databaseDir = argv[3];   //训练用数据集
    string resultDir = argv[4]; //词典存放位置
    if ( argc == 8 ) {
        params.detectorType = argv[5];
        params.descriptorType = argv[6];
        params.matcherType = argv[7];
    }
    resultDir = QCoreApplication::applicationDirPath().toStdString() + "/" + resultDir;
    databaseDir = QCoreApplication::applicationDirPath().toStdString() + "/" + databaseDir;
    MakeDir(resultDir);

    // key: image category name
    // value: histogram of image
    vector<string> abstractCategories;
    GetDirList(databaseDir, abstractCategories);    //抽象类别
    //    vector<map<string, Mat> > absSamples;  //保存samples的中间结果, 类名-图像描述符

    for(int i = 0; i < (int)abstractCategories.size(); ++i) {
        Mat inheritedVocabulary;
        string currentAbsDatabaseDir = databaseDir + "/" + abstractCategories[i];
        string currentResultDir = resultDir + "/" + abstractCategories[i];
        MakeDir(currentResultDir);
        MakeDir(currentResultDir + "/BagofWords");
        MakeDir(currentResultDir + "/svm");
        vector<string> inheritedCategories;
        GetDirList(currentAbsDatabaseDir,inheritedCategories);
        BOOST_AUTO(samples,
                   bow(params,
                       currentResultDir + "/" + kVocabularyFile,    //词典所在目录
                       inheritedVocabulary,
                       currentAbsDatabaseDir,
                       currentResultDir + "/BagofWords",    //BoW文件所在目录
                       inheritedCategories));
        trainSVM(samples, currentResultDir + "/svm");
        MakeDir(currentResultDir + "/SelectedFiles");
        MakeDir(currentResultDir + "/SelectedFiles/AbstractCategorySVM");
        MakeDir(currentResultDir + "/SelectedFiles/AbstractBagofWords");
        randomSelectImage(databaseDir, abstractCategories[i], currentResultDir + "/SelectedFiles");    //从其他抽象类中选取部分图像用于训练当前抽象类的词典
        Mat absVocabulary;
        vector<string> absCategories;
        GetDirList(currentResultDir + "/SelectedFiles",absCategories);
        BOOST_AUTO(pos, bow(params,
                            currentResultDir + "/SelectedFiles/" + kVocabularyFile,
                            absVocabulary,
                            currentResultDir + "/SelectedFiles",
                            currentResultDir + "/SelectedFiles/AbstractBagofWords",
                            absCategories));
        trainSVM(pos, currentResultDir + "/SelectedFiles/AbstractCategorySVM");
    }
    cout<<"###I### Done processing abstracted categories."<<endl;
    cout<<"Starting classification test..."<<endl;
    test(params,testDir);
    return 0;
}

void help( const char* progName ) {
    cout << "Usage: \n"
         << progName << " {classify method} {query image} {image set path} {result directory}\n"
         << "or: \n"
         << progName << " {classify method} {image set path} {result directory} {feature detector} {descriptor extrator} {descriptor matcher}\n"
         << "\n"
         << "Input parameters: \n"
         << "{classify method}			\n	Method used to classify image, can be one of svm or match.\n"
         << "{query image}				\n	Path to query image.\n"
         << "{image set path}			\n	Path to image training set, organized into categories, like Caltech 101.\n"
         << "{result directory}			\n	Path to result directory.\n"
         << "{feature detector}			\n	Feature detector name, should be one of\n"
         <<	"	FAST, STAR, SIFT, SURF, MSER, GFTT, HARRIS.\n"
         << "{descriptor extractor}		\n	Descriptor extractor name, should be one of \n"
         <<	"	SURF, OpponentSIFT, SIFT, OpponentSURF, BRIEF.\n"
         << "{descriptor matcher}		\n	Descriptor matcher name, should be one of\n"
         << "	BruteForce, BruteForce-L1, FlannBased, BruteForce-Hamming, BruteForce-HammingLUT.\n";
}

bool MakeDir( const string& dirAbsPath ) {
    QDir dir;
    if(dir.mkdir(QString::fromStdString(dirAbsPath)) || dir.exists())
        return true;
    cout<<"###E### can't make dir!"<<dirAbsPath<<endl;
    return false;
}

void List(string sourceDir, vector<string> &dirNames, QDir::Filters filter) {
    QDir dir(QString::fromStdString(sourceDir));
    if(!dir.exists()) {
        cout<<"###E### dir does not exist"<<endl;
        return;
    }
    dir.setFilter(filter | QDir::Hidden | QDir::NoSymLinks | QDir::NoDotAndDotDot);
    dir.setSorting(QDir::Size | QDir::Reversed);

    QFileInfoList list = dir.entryInfoList();
    for(int i = 0; i < list.size(); ++i) {
        QFileInfo fileInfo = list.at(i);
        //        cout<<"name="<<fileInfo.fileName().toStdString()<<endl;
        dirNames.push_back(fileInfo.fileName().toStdString());
    }
}
//列出sourceDir中所有目录的名称，存放于dirNames中
void GetDirList(string sourceDir, vector<string> &dirNames) {
    List(sourceDir,dirNames,QDir::Dirs);
}

void GetFileList(string sourceDir, vector<string>& fileList ) {
    List(sourceDir,fileList,QDir::Files);
}
